#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
from PIL import Image
import numpy as np

def list_files_in_directory(directory):
    """
    遍历指定文件夹下的所有文件名

    参数:
    directory (str): 要遍历的文件夹路径

    返回:
    list: 文件名列表
    """
    file_names = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_names.append(file)
    return file_names

def resize_crop_padding(image_path, new_size):
    """
    加载图像，首先resize，然后crop，最后padding为new size的尺寸。

    参数:
    image_path (str): 图像文件的路径
    new_size (tuple): 新的尺寸，格式为 (width, height)

    返回:
    list: new_size的图像对象列表
    """
    # 打开图像文件
    with Image.open(image_path) as img:
        # 转换为灰度模式
        img_gray = img.convert('L')
        # 获取图像的尺寸
        width, height = img_gray.size
        
        # 计算缩放比例
        scale = float(new_size[0] * 2) / width
        
        # 调整图像大小
        resized_img = img_gray.resize((int(new_size[0] * 2), int(height * scale)), Image.LANCZOS)
        
        # 获取调整后的图像尺寸
        resized_width, resized_height = resized_img.size
        
        # 计算分割后的子图像尺寸
        sub_height = resized_height // 2
        sub_width = resized_width // 2
        
        # 分割图像为四份
        subImg1 = resized_img.crop((0, 0, sub_width, sub_height))
        subImg2 = resized_img.crop((sub_width, 0, resized_width, sub_height))
        subImg3 = resized_img.crop((0, sub_height, sub_width, resized_height))
        subImg4 = resized_img.crop((sub_width, sub_height, resized_width, resized_height))
        
        # 将图像转换为numpy数组以便于填充
        subImg1_np = np.array(subImg1)
        subImg2_np = np.array(subImg2)
        subImg3_np = np.array(subImg3)
        subImg4_np = np.array(subImg4)
        
        # 计算需要补零的高度和宽度
        pad_height = new_size[1] - sub_height
        pad_width = new_size[0] - sub_width
        
        # 计算上下左右补零的大小
        pad_top = pad_height // 2
        pad_bottom = pad_height - pad_top
        pad_left = pad_width // 2
        pad_right = pad_width - pad_left
        
        # 使用numpy的pad函数在图像的上下左右方向补零
        subImg1_padded = np.pad(subImg1_np, ((pad_top, pad_bottom), (pad_left, pad_right)), 'constant', constant_values=0)
        subImg2_padded = np.pad(subImg2_np, ((pad_top, pad_bottom), (pad_left, pad_right)), 'constant', constant_values=0)
        subImg3_padded = np.pad(subImg3_np, ((pad_top, pad_bottom), (pad_left, pad_right)), 'constant', constant_values=0)
        subImg4_padded = np.pad(subImg4_np, ((pad_top, pad_bottom), (pad_left, pad_right)), 'constant', constant_values=0)
        
        # 将填充后的图像转换回PIL图像
        subImg1_padded_img = Image.fromarray(subImg1_padded)
        subImg2_padded_img = Image.fromarray(subImg2_padded)
        subImg3_padded_img = Image.fromarray(subImg3_padded)
        subImg4_padded_img = Image.fromarray(subImg4_padded)
        
        # 将填充后的图像存储在列表中
        resized_cropped_padded_images = [subImg1_padded_img, subImg2_padded_img, subImg3_padded_img, subImg4_padded_img]
        
        return resized_cropped_padded_images

# 此函数用于测试，跟windows画图函数的效果进行比较
def resize_image(image_path, new_size):
    """
    加载图像，resize为new_size

    参数:
    image_path (str): 图像文件的路径
    new_size (tuple): 新的尺寸，格式为 (width, height)

    返回:
    缩放后的图像对象
    """
    # 打开图像文件
    with Image.open(image_path) as img:
        # 转换为灰度模式
        img_gray = img.convert('L')
            
        resized_image = img_gray.resize(new_size, Image.LANCZOS)

        return resized_image


if __name__ == '__main__':
    # 示例用法
    directory_path = 'E:/YCS_Temp/project/pabone/thesis_code/simulation/STARE_Dataset/training'
    new_size = (268, 534)  # 新的尺寸为268x534
    #new_size = (536, 464)
    new_directory_path = 'E:/YCS_Temp/project/pabone/thesis_code/simulation/VESSEL_Dataset'

    files = list_files_in_directory(directory_path)
    for file in files:
        file_path = os.path.join(directory_path, file)
        print(file_path)

        resized_cropped_padded_images  = resize_crop_padding(file_path, new_size)

        # # 显示结果
        # for i, img in enumerate(resized_cropped_padded_images):
        #     img.show(title=f'Padded Image {i+1}')

        ## 保存裁剪并缩放后的图像
        for i, img in enumerate(resized_cropped_padded_images):
           img.save(os.path.join(new_directory_path, file + f'_{i}.png'))

        # ## 仅用于测试，跟windows画图函数的效果进行比较
        # resized_img = resize_image(file_path, new_size)
        # resized_img.save(os.path.join(new_directory_path, file + f'_lanczos.png'))
        
        #break